#!/usr/bin/env python3
"""
自提示机制模块

实现自提示机制，让模型自动生成初始提示并进行迭代优化
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple, Optional, Any
import cv2

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from .base_prompt import BasePromptGenerator
from utils.logger import get_logger


class SelfPromptGenerator(BasePromptGenerator):
    """自提示生成器"""
    
    def __init__(self, 
                 input_dim: int = 256,
                 prompt_dim: int = 256,
                 max_iterations: int = 3,
                 confidence_threshold: float = 0.7,
                 device: str = "cuda"):
        """
        初始化自提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            max_iterations: 最大迭代次数
            confidence_threshold: 置信度阈值
            device: 设备
        """
        super().__init__(input_dim, prompt_dim, device)
        
        self.max_iterations = max_iterations
        self.confidence_threshold = confidence_threshold
        
        # 初始提示生成器
        self.initial_prompt_generator = nn.Sequential(
            nn.Linear(input_dim, prompt_dim),
            nn.ReLU(),
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 2 + 4 + 32 * 32)  # points + boxes + masks
        ).to(device)
        
        # 提示优化器
        self.prompt_optimizer = nn.Sequential(
            nn.Linear(prompt_dim + 2 + 4 + 32 * 32, prompt_dim),  # 特征 + 当前提示
            nn.ReLU(),
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 2 + 4 + 32 * 32)  # 优化后的提示
        ).to(device)
        
        # 置信度评估器
        self.confidence_evaluator = nn.Sequential(
            nn.Linear(prompt_dim + 2 + 4 + 32 * 32, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 1),
            nn.Sigmoid()
        ).to(device)
        
        # 提示质量评估器
        self.quality_evaluator = nn.Sequential(
            nn.Linear(2 + 4 + 32 * 32, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        ).to(device)
        
        # 注意力机制用于提示融合
        self.attention = nn.MultiheadAttention(
            embed_dim=prompt_dim,
            num_heads=8,
            dropout=0.1,
            batch_first=True
        ).to(device)
        
        # 提示历史记录
        self.prompt_history = []
        
        self.logger.info(f"自提示生成器初始化完成: {self.get_model_info()}")
    
    def generate_initial_prompts(self, image_features: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        生成初始提示
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            
        Returns:
            初始提示字典
        """
        batch_size = image_features.shape[0]
        
        # 全局平均池化
        global_features = F.adaptive_avg_pool2d(image_features, (1, 1)).flatten(1)  # (B, C)
        
        # 确保特征维度匹配
        if global_features.shape[1] != self.input_dim:
            # 如果特征维度不匹配，使用投影层
            if not hasattr(self, 'feature_projection'):
                self.feature_projection = nn.Linear(global_features.shape[1], self.input_dim).to(self.device)
            global_features = self.feature_projection(global_features)
        
        # 生成初始提示
        initial_prompts_raw = self.initial_prompt_generator(global_features)  # (B, 2+4+32*32)
        
        # 分离不同类型的提示
        points = initial_prompts_raw[:, :2].unsqueeze(1)  # (B, 1, 2)
        boxes = initial_prompts_raw[:, 2:6].unsqueeze(1)  # (B, 1, 4)
        masks = initial_prompts_raw[:, 6:].view(batch_size, 1, 32, 32)  # (B, 1, 32, 32)
        
        # 归一化坐标到[0, 1]，并确保有合理的范围
        points = torch.sigmoid(points)
        boxes = torch.sigmoid(boxes)
        masks = torch.sigmoid(masks)
        
        # 确保提示有合理的分布
        # 将points和boxes缩放到更大的范围
        points = points * 0.8 + 0.1  # 范围[0.1, 0.9]
        boxes = boxes * 0.8 + 0.1   # 范围[0.1, 0.9]
        
        return {
            'points': points,
            'boxes': boxes,
            'masks': masks,
            'iteration': 0
        }
    
    def evaluate_prompt_quality(self, prompts: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        评估提示质量
        
        Args:
            prompts: 提示字典
            
        Returns:
            质量分数 (B, 1)
        """
        # 拼接所有提示
        points = prompts['points'].view(prompts['points'].size(0), -1)  # (B, 2)
        boxes = prompts['boxes'].view(prompts['boxes'].size(0), -1)    # (B, 4)
        masks = prompts['masks'].view(prompts['masks'].size(0), -1)    # (B, 32*32)
        
        all_prompts = torch.cat([points, boxes, masks], dim=1)  # (B, 2+4+32*32)
        
        # 评估质量
        quality_score = self.quality_evaluator(all_prompts)  # (B, 1)
        
        return quality_score
    
    def optimize_prompts(self, 
                        image_features: torch.Tensor,
                        current_prompts: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        优化提示
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            current_prompts: 当前提示
            
        Returns:
            优化后的提示
        """
        batch_size = image_features.shape[0]
        
        # 提取特征
        global_features = F.adaptive_avg_pool2d(image_features, (1, 1)).flatten(1)  # (B, C)
        
        # 确保特征维度匹配
        if global_features.shape[1] != self.input_dim:
            if not hasattr(self, 'feature_projection'):
                self.feature_projection = nn.Linear(global_features.shape[1], self.input_dim).to(self.device)
            global_features = self.feature_projection(global_features)
        
        # 拼接当前提示
        points = current_prompts['points'].view(batch_size, -1)  # (B, 2)
        boxes = current_prompts['boxes'].view(batch_size, -1)    # (B, 4)
        masks = current_prompts['masks'].view(batch_size, -1)    # (B, 32*32)
        
        current_prompts_flat = torch.cat([points, boxes, masks], dim=1)  # (B, 2+4+32*32)
        
        # 融合特征和提示
        combined_input = torch.cat([global_features, current_prompts_flat], dim=1)  # (B, C+2+4+32*32)
        
        # 优化提示
        optimized_prompts_raw = self.prompt_optimizer(combined_input)  # (B, 2+4+32*32)
        
        # 分离优化后的提示
        opt_points = optimized_prompts_raw[:, :2].unsqueeze(1)  # (B, 1, 2)
        opt_boxes = optimized_prompts_raw[:, 2:6].unsqueeze(1)  # (B, 1, 4)
        opt_masks = optimized_prompts_raw[:, 6:].view(batch_size, 1, 32, 32)  # (B, 1, 32, 32)
        
        # 归一化
        opt_points = torch.sigmoid(opt_points)
        opt_boxes = torch.sigmoid(opt_boxes)
        opt_masks = torch.sigmoid(opt_masks)
        
        return {
            'points': opt_points,
            'boxes': opt_boxes,
            'masks': opt_masks
        }
    
    def compute_confidence(self, 
                          image_features: torch.Tensor,
                          prompts: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算提示置信度
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            prompts: 提示字典
            
        Returns:
            置信度 (B, 1)
        """
        batch_size = image_features.shape[0]
        
        # 提取特征
        global_features = F.adaptive_avg_pool2d(image_features, (1, 1)).flatten(1)  # (B, C)
        
        # 确保特征维度匹配
        if global_features.shape[1] != self.input_dim:
            if not hasattr(self, 'feature_projection'):
                self.feature_projection = nn.Linear(global_features.shape[1], self.input_dim).to(self.device)
            global_features = self.feature_projection(global_features)
        
        # 拼接提示
        points = prompts['points'].view(batch_size, -1)  # (B, 2)
        boxes = prompts['boxes'].view(batch_size, -1)    # (B, 4)
        masks = prompts['masks'].view(batch_size, -1)    # (B, 32*32)
        
        prompts_flat = torch.cat([points, boxes, masks], dim=1)  # (B, 2+4+32*32)
        
        # 融合特征和提示
        combined_input = torch.cat([global_features, prompts_flat], dim=1)  # (B, C+2+4+32*32)
        
        # 计算置信度
        confidence = self.confidence_evaluator(combined_input)  # (B, 1)
        
        return confidence
    
    def iterative_optimization(self, 
                             image_features: torch.Tensor) -> Dict[str, Any]:
        """
        迭代优化提示
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            
        Returns:
            优化结果字典
        """
        # 生成初始提示
        current_prompts = self.generate_initial_prompts(image_features)
        
        # 记录历史
        self.prompt_history = [current_prompts.copy()]
        
        best_prompts = current_prompts.copy()
        best_confidence = torch.zeros(image_features.size(0), 1, device=self.device)
        
        for iteration in range(1, self.max_iterations + 1):
            # 优化提示
            optimized_prompts = self.optimize_prompts(image_features, current_prompts)
            optimized_prompts['iteration'] = iteration
            
            # 计算置信度
            confidence = self.compute_confidence(image_features, optimized_prompts)
            
            # 计算质量分数
            quality = self.evaluate_prompt_quality(optimized_prompts)
            
            # 更新最佳提示
            for i in range(image_features.size(0)):
                if confidence[i] > best_confidence[i]:
                    best_confidence[i] = confidence[i]
                    best_prompts['points'][i] = optimized_prompts['points'][i]
                    best_prompts['boxes'][i] = optimized_prompts['boxes'][i]
                    best_prompts['masks'][i] = optimized_prompts['masks'][i]
            
            # 记录历史
            self.prompt_history.append(optimized_prompts.copy())
            
            # 检查收敛条件
            if torch.all(confidence > self.confidence_threshold):
                self.logger.info(f"在第 {iteration} 次迭代后收敛")
                break
            
            # 更新当前提示
            current_prompts = optimized_prompts
        
        # 添加最终信息
        best_prompts['confidence'] = best_confidence
        best_prompts['quality'] = self.evaluate_prompt_quality(best_prompts)
        best_prompts['total_iterations'] = len(self.prompt_history)
        
        return best_prompts
    
    def forward(self, 
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            提示字典
        """
        # 迭代优化
        result = self.iterative_optimization(image_features)
        
        # 只返回标准提示格式，不包含额外信息
        return {
            'points': result['points'],
            'boxes': result['boxes'],
            'masks': result['masks']
        }
    
    def generate_prompts(self, 
                        image: torch.Tensor,
                        metadata: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        生成提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            metadata: 元数据
            
        Returns:
            提示字典
        """
        # 提取图像特征
        image_features = self.extract_features(image)
        
        # 生成自提示
        prompts = self.forward(image_features, metadata)
        
        return prompts
    
    def compute_self_prompt_loss(self, 
                               predictions: Dict[str, torch.Tensor],
                               targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算自提示损失
        
        Args:
            predictions: 预测结果
            targets: 目标结果
            
        Returns:
            损失值
        """
        total_loss = 0.0
        
        # 基础提示损失
        total_loss += self.compute_loss(predictions, targets)
        
        # 置信度损失
        if 'confidence' in predictions and 'confidence' in targets:
            confidence_loss = F.mse_loss(
                predictions['confidence'],
                targets['confidence']
            )
            total_loss += 0.1 * confidence_loss
        
        # 质量损失
        if 'quality' in predictions and 'quality' in targets:
            quality_loss = F.mse_loss(
                predictions['quality'],
                targets['quality']
            )
            total_loss += 0.1 * quality_loss
        
        # 迭代一致性损失
        if len(self.prompt_history) > 1:
            consistency_loss = 0.0
            for i in range(1, len(self.prompt_history)):
                prev_prompts = self.prompt_history[i-1]
                curr_prompts = self.prompt_history[i]
                
                # 计算相邻迭代间的差异
                points_diff = F.mse_loss(prev_prompts['points'], curr_prompts['points'])
                boxes_diff = F.mse_loss(prev_prompts['boxes'], curr_prompts['boxes'])
                masks_diff = F.mse_loss(prev_prompts['masks'], curr_prompts['masks'])
                
                consistency_loss += points_diff + boxes_diff + masks_diff
            
            consistency_loss /= (len(self.prompt_history) - 1)
            total_loss += 0.05 * consistency_loss
        
        return total_loss
    
    def get_optimization_history(self) -> List[Dict[str, torch.Tensor]]:
        """
        获取优化历史
        
        Returns:
            优化历史列表
        """
        return self.prompt_history.copy()
    
    def clear_history(self):
        """清空优化历史"""
        self.prompt_history.clear()
        self.logger.info("优化历史已清空")
    
    def get_optimization_stats(self) -> Dict[str, Any]:
        """
        获取优化统计信息
        
        Returns:
            统计信息字典
        """
        if not self.prompt_history:
            return {
                'total_iterations': 0,
                'converged': False,
                'final_confidence': 0.0,
                'final_quality': 0.0
            }
        
        final_prompts = self.prompt_history[-1]
        
        return {
            'total_iterations': len(self.prompt_history),
            'converged': len(self.prompt_history) < self.max_iterations,
            'final_confidence': final_prompts.get('confidence', torch.tensor(0.0)).mean().item(),
            'final_quality': final_prompts.get('quality', torch.tensor(0.0)).mean().item()
        }
